package com.tuniu.agents.impl.trip;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.common.entity.trip.TripSearchResponse;
import com.tuniu.agents.common.handler.IStreamResponseHandler;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

@Slf4j
public class SearchStreamResponseHandler implements IStreamResponseHandler {
    public static final String CONVERSATION_ID = "conversationId";
    public static final String VECTOR_MATCH_URL = "vectorMatchUrl";
    public static final String HEADERS = "headers";
    public static final String TRIP_SUMMARY = "tripSummary";
    public static final Integer TOP_K = 3;
    public static final String TRIP_SEARCH_KEY = "MATCHED_TRIP_SUMMARY";
    private final ChatEnvironmentMemory chatEnvironmentMemory;
    private AgentManager agentManager;

    public SearchStreamResponseHandler(ChatEnvironmentMemory chatEnvironmentMemory, AgentManager agentManager) {
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.agentManager = agentManager;
    }


    @Override
    public Flux<String> handle(Map<String, Object> handlerContext, Flux<String> flux) {
        String conversationId = (String) handlerContext.get(CONVERSATION_ID);
        String vectorMatchUrl = (String) handlerContext.get(VECTOR_MATCH_URL);
        Map<String, String> headers = (Map<String, String>) handlerContext.get(HEADERS);
        String tripSummary = (String) handlerContext.get(TRIP_SUMMARY);
        log.info("SearchStreamResponseHandler processing for trip summary: {}  for conversation: {}", tripSummary, conversationId);
        
        if (tripSummary == null || vectorMatchUrl == null) {
            log.error("Missing required parameters for vector matching. tripSummary: {}, vectorMatchUrl: {}", 
                      tripSummary, vectorMatchUrl);
            return Flux.just("Unable to process your trip summary, please provide more detailed information.");
        }
        
        Map<String, Object> requestVectorMatch = new HashMap<>();
        requestVectorMatch.put("query_text", tripSummary);
        requestVectorMatch.put("top_k", TOP_K);
        
        OkHttpClient client = new OkHttpClient();
        
        try {
            log.info("Making vector match request to URL: {} with request: {}", vectorMatchUrl, requestVectorMatch);
            
            MediaType JSON = MediaType.parse("application/json; charset=utf-8");
            String requestJson = new ObjectMapper().writeValueAsString(requestVectorMatch);
            RequestBody requestBody = RequestBody.create(JSON, requestJson);
            
            Request.Builder requestBuilder = new Request.Builder()
                    .url(vectorMatchUrl)
                    .post(requestBody);

            // Add headers if available
            if (headers != null) {
                for (Map.Entry<String, String> entry : headers.entrySet()) {
                    requestBuilder.addHeader(entry.getKey(), entry.getValue());
                }
            }

            Request request = requestBuilder.build();
            Response okHttpResponse = client.newCall(request).execute();
            String response = okHttpResponse.body().string().replace("\\", "").replace("\"{", "{").replace("}\"", "}");

            log.info("Vector match response: {}", response);

            if (response == null || response.isEmpty()) {
                log.warn("Empty response from vector matching service");
                return Flux.just("No information found related to your trip, please try providing a more detailed trip description.");
            }

            return processVectorMatchResponse(response, conversationId);
        } catch (Exception e) {
            log.error("Error during vector matching request", e);
            return Flux.just("Search service is temporarily unavailable, please try again later.");
        }
    }


    private Flux<String> processVectorMatchResponse(String response, String conversationId) {
        try {
            TripSearchResponse tripSearchResponse = JsonUtil.fromJson(response, TripSearchResponse.class);
            String productIds = Optional.ofNullable(tripSearchResponse.getData())
                .map(dataList -> dataList.stream()
                    .map(data -> String.valueOf(data.getProduct_id()))
                    .collect(Collectors.joining(",")))
                .orElse("");
            Map<String, Object> responseMap = new HashMap<>();
            responseMap.put(TRIP_SEARCH_KEY, tripSearchResponse);
            chatEnvironmentMemory.add(conversationId, responseMap);


            return Flux.just(tripSearchResponse.toMarkdown());

            // Agent agent = agentManager.getAgent("ItineraryPlanner");
            // return agent.stream(conversationId, new String());

        } catch (Exception e) {
            log.error("Error parsing vector match response for conversation: {}", conversationId, e);
            return Flux.just("Error occurred while parsing search results, please try again later.");
        }
    }
}
